import random
import socket
import struct
from enum import Enum
from typing import Self
import math

import libs.hashs as hashs

HOST = 'www.baidu.com'
PORT = 443
TLS_VERSION = b'\x03\x01'
TLS_RSA_WITH_RC4_128_SHA = b'\x00\x05'

class ContentType(Enum):
    change_cipher_spec = bytes([20])
    alert = bytes([21])
    handshake = bytes([22])
    application_data = bytes([23])
    
    @staticmethod
    def from_int(value: int):
        return ContentType(bytes([value]))


class HandshakeType(Enum):
    hello_request = bytes([0])
    client_hello = bytes([1])
    server_hello = bytes([2])
    certificate = bytes([11])
    server_key_exchange = bytes([12])
    certificate_request = bytes([13])
    server_hello_done = bytes([14])
    certificate_verify = bytes([15])
    client_key_exchange = bytes([16])
    finished = bytes([20])
    
    @staticmethod
    def from_int(value: int):
        return HandshakeType(bytes([value]))


class AlertLevel(Enum):
    warning = 1
    fatal = 2


class AlertDescription(Enum):
    close_notify = 0
    unexpected_message = 10
    bad_record_mac = 20
    decryption_failed = 21
    record_overflow = 22
    decompression_failure = 30
    handshake_failure = 40
    bad_certificate = 42
    unsupported_certificate = 43
    certificate_revoked = 44
    certificate_expired = 45
    certificate_unknown = 46
    illegal_parameter = 47
    unknown_ca = 48
    access_denied = 49
    decode_error = 50
    decrypt_error = 51
    export_restriction = 60
    protocol_version = 70
    insufficient_security = 71
    internal_error = 80
    user_canceled = 90
    no_renegotiation = 100

def TLSRecordLayer(contentType: ContentType, payload: bytes):
    return contentType.value + TLS_VERSION + struct.pack('>H', len(payload)) + payload

def P_HASH(secret: bytes, seed: bytes, length: int, hash_func: hashs.hashcls):
    A = seed
    p = b""
    while len(p) <= length:
        A = hashs.HMAC(secret, A, hash_func)
        p += hashs.HMAC(secret, A+seed, hash_func)

    return p[0:length]

def PRF(secret: bytes, label: bytes, seed: bytes, length: int):
    mid = int(math.ceil(len(secret)/2))
    S1 = secret[:mid]
    S2 = secret[-mid:]

    pmd5 = P_HASH(S1, label+seed, length, hashs.MD5)
    psha1 = P_HASH(S2, label+seed, length, hashs.SHA1)
    return bytes(a ^ b for a, b in zip(pmd5, psha1))

class TLS:
    def __init__(self: Self):

        self.client_random = random.randbytes(32)
        self.server_random = b''

        self.md5 = hashs.MD5()
        self.sha1 = hashs.SHA1()

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)

        self.keepAlive = True
        self.serverCCS = False
        self.serverHelloDone = False
        self.serverFinished = False

        self.sequence_number_client = 0
        self.sequence_number_server = 0

        self.session_keys_client_rc4 = None
        self.session_keys_server_rc4 = None
        self.session_keys_server_pubkey = None
        self.session_keys_master = None
        self.session_keys_client_write_MAC_secret = None
        self.session_keys_server_write_MAC_secret = None
        self.session_keys_client_write_IV = None
        self.session_keys_server_write_IV = None

    def run(self: Self):
        self.sock.connect((HOST, PORT))
        self.client_hello()
        while not self.serverHelloDone:
            # ServerHello
            # Certificate
            # ServerHelloDone
            self.recv_from_server()
        self.client_key_exchange()
        self.client_change_cipher_spec()
        self.client_finished()
        while not self.serverFinished:
            # ChangeCipherSpec
            # Finished
            self.recv_from_server()
        
        self.send_app_data('\r\n'.join([
            'GET / HTTP/1.1',
            f'HOST: {HOST}',
            'Accept-Encoding: identity',
            'Connection: close',
        ]) + '\r\n\r\n')
        self.sock.shutdown(socket.SHUT_WR)
        s = b''
        while self.keepAlive:
            r = self.recv_from_server()
            if isinstance(r, bytes):
                s += r
        print(s.decode('utf-8'))
        self.sock.shutdown(socket.SHUT_RD)
        self.sock.close()
        print()

    def client_hello(self: Self):
        hello = TLS_VERSION + self.client_random
        hello += b'\x00' # session id length
        hello += b'\x00\x02' # cipher suite length
        hello += TLS_RSA_WITH_RC4_128_SHA
        hello += b'\x01\x00' # Compression method NULL
        buf = HandshakeType.client_hello.value + b'\x00' + struct.pack('>H', len(hello)) + hello
        self.md5.update(buf)
        self.sha1.update(buf)
        req = TLSRecordLayer(ContentType.handshake, buf)
        self.sock.send(req)

    def client_key_exchange(self: Self):
        premaster_secret = TLS_VERSION + random.randbytes(46)
        cipherpremaster = self.session_keys_server_pubkey.encrypt(premaster_secret)

        clientkeyexchange = HandshakeType.client_key_exchange.value
        clientkeyexchange += b"\x00" + struct.pack(">H", len(cipherpremaster)+2)
        clientkeyexchange += struct.pack(">H",len(cipherpremaster))
        clientkeyexchange += cipherpremaster
        self.md5.update(clientkeyexchange)
        self.sha1.update(clientkeyexchange)
        record_cke = TLSRecordLayer(ContentType.handshake, clientkeyexchange)
        self.sock.send(record_cke)

        # Calculate master secret
        self.session_keys_master = PRF(premaster_secret, b"master secret", self.client_random+self.server_random, 48)

        # Calculate keys, MAC secret and IV
        outputLength = 20*2 + 16*2  # MAC size 20 + Key size 16
        keyBlock = PRF(self.session_keys_master, b"key expansion", self.server_random+self.client_random, outputLength)

        self.session_keys_client_write_MAC_secret = keyBlock[0:20]
        self.session_keys_server_write_MAC_secret = keyBlock[20:40]
        session_keys_client_write_key = keyBlock[40:56]
        session_keys_server_write_key = keyBlock[56:72]
        self.session_keys_client_rc4 = hashs.ARC4(session_keys_client_write_key)
        self.session_keys_server_rc4 = hashs.ARC4(session_keys_server_write_key)

    def client_change_cipher_spec(self: Self):
        record_ccs = TLSRecordLayer(ContentType.change_cipher_spec, b"\x01")
        self.sock.send(record_ccs)

    def client_finished(self: Self):
        # Create Encrypted Finished message
        clientfinished = HandshakeType.finished.value
        clientfinished += b"\x00\x00\x0C" # Length = 12 bytes
        hs_hash = self.md5.digest()+self.sha1.digest()
        verify = PRF(self.session_keys_master, b"client finished", hs_hash, 12)
        clientfinished += verify
        self.md5.update(clientfinished)
        self.sha1.update(clientfinished)

        self.send_encrypted_data(ContentType.handshake, clientfinished)

    def send_app_data(self: Self, data: str):
        self.send_encrypted_data(ContentType.application_data, data.encode('utf-8'))

    def send_encrypted_data(self: Self, contentType: ContentType, data: bytes):
        record = TLSRecordLayer(contentType, data)

        # Calculate HMAC
        seqnumber = struct.pack('!Q', self.sequence_number_client)
        self.sequence_number_client += 1
        recordmac = hashs.HMAC(self.session_keys_client_write_MAC_secret, seqnumber+record, hashs.SHA1)
        tobeencrypted = data + recordmac

        ciphertext = self.session_keys_client_rc4.encrypt(tobeencrypted)
        self.sock.send(TLSRecordLayer(contentType, ciphertext))

    def recv_from_server(self: Self):
        r = self.sock.recv(5)
        if len(r) < 1:
            self.keepAlive = False
            return
        
        contentType, v, n = struct.unpack('>B2H', r)
        assert v == 0x301
        contentType = ContentType.from_int(contentType)
        buf = b''
        while len(buf) < n:
            buf += self.sock.recv(n - len(buf))
        if self.serverCCS:
            buf = self.decodeData(buf, contentType)
        match contentType:
            case ContentType.change_cipher_spec:
                self.serverCCS = True
            case ContentType.alert:
                assert len(buf) == 2
                level, desc = AlertLevel(buf[0]), AlertDescription(buf[1])
                match level:
                    case AlertLevel.warning:
                        match desc:
                            case AlertDescription.close_notify:
                                self.keepAlive = False
                            case _:
                                print(f'warning: {desc.name}')
                    case AlertLevel.fatal:
                        raise Exception(f'fatal: {desc.name}')
            case ContentType.handshake:
                self.ExtractHelloProto(buf)
            case ContentType.application_data:
                return buf

    def decodeData(self: Self, payload: bytes, contentType: ContentType):
        plaintext = self.session_keys_server_rc4.decrypt(payload)

        # get MAC
        servermac = plaintext[-20:]
        plaintext = plaintext[:-20] # strip MAC

        seqnumber = struct.pack('!Q', self.sequence_number_server) 
        self.sequence_number_server += 1
        record = TLSRecordLayer(contentType, plaintext)
        recordmac = hashs.HMAC(self.session_keys_server_write_MAC_secret, seqnumber+record, hashs.SHA1)

        if recordmac != servermac:
            raise Exception(f"HMAC INVALID: {recordmac} != {servermac}")

        return plaintext
    
    def ExtractHelloProto(self: Self, buf: bytes):
        handshakeType = HandshakeType.from_int(buf[0])
        n, = struct.unpack('>L', b'\x00' + buf[1:4])
        sub_buf = buf[4:]
        assert n == len(sub_buf), f'{n} != {len(sub_buf)}'
        if handshakeType != HandshakeType.finished:
            self.md5.update(buf)
            self.sha1.update(buf)
        match handshakeType:
            case HandshakeType.server_hello:
                self.server_random = sub_buf[2:34]
            case HandshakeType.certificate:
                certlen, = struct.unpack('>L', b'\x00' + sub_buf[3:6])
                certbuf = sub_buf[6:6+certlen]
                cert = hashs.der_decode(certbuf)[0]
                cert = hashs.der_decode(cert)[6]
                self.session_keys_server_pubkey = hashs.RSA.from_der(cert)
            case HandshakeType.server_hello_done:
                self.serverHelloDone = True
            case HandshakeType.finished:
                verify = PRF(self.session_keys_master, b"server finished", self.md5.digest()+self.sha1.digest(),12)
                if verify != sub_buf:
                    raise Exception(f'verify({verify}) != serververify({sub_buf})')
                self.serverFinished = True
            case _:
                raise Exception(f'not Implemented type: {handshakeType.name}')

if __name__=='__main__':
    TLS().run()